# coding=utf-8
# 全世界最帅的帅哥
# 测试链接: https://leetcode.cn/problems/partition-list/description/
from typing import Optional

from utils.list2linked_list import list2linked_list

head = list2linked_list([1])


class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next


class Solution:
    def partition(self, head: Optional[ListNode], x: int) -> Optional[ListNode]:
        small_head, big_head = None, None
        small_node, big_node = None, None

        while head:
            if head.val < x:
                if small_head:
                    small_head.next = ListNode(head.val)
                    small_head = small_head.next
                else:
                    small_head = ListNode(head.val)
                    small_node = small_head
            else:
                if big_head:
                    big_head.next = ListNode(head.val)
                    big_head = big_head.next
                else:
                    big_head = ListNode(head.val)
                    big_node = big_head

            head = head.next
        if not small_head:
            return big_node
        small_head.next = big_node
        return small_node


def test():
    s = Solution()
    s.partition(head, 1)


if __name__ == '__main__':
    test()
